import os

current_path = os.path.dirname(os.path.abspath(__file__))
os.environ["PATH"] = os.path.join(current_path, "ffmpeg/bin/") + ";" + os.environ["PATH"]
import torch
import commons
import utils
import re
from models import SynthesizerTrn
from text import text_to_sequence_with_hps as text_to_sequence
from scipy.io.wavfile import write
from pydub import AudioSegment

device = torch.device("cpu")  # cpu  mps
hps = utils.get_hparams_from_file("{}/configs/finetune_speaker.json".format(current_path))
hps.data.text_cleaners[0] = 'my_infer_ce_cleaners'
hps.data.n_speakers = 2
symbols = hps.symbols
net_g = SynthesizerTrn(
    len(symbols),
    hps.data.filter_length // 2 + 1,
    hps.train.segment_size // hps.data.hop_length,
    n_speakers=hps.data.n_speakers,
    **hps.model).to(device)
_ = net_g.eval()
#    G_latest  G_trilingual  G_930000  G_953000 G_984000 G_990000 G_1004000 G_1021000
_ = utils.load_checkpoint("{}/OUTPUT_MODEL/G_1021000.pth".format(current_path), net_g, None)


def add_laug_tag(text):
    '''
    添加语言标签
    '''
    pattern = r'([\u4e00-\u9fa5，。！？；：、——……（）]+|[a-zA-Z,.:()]+|\d+)'
    segments = re.findall(pattern, text)
    for i in range(len(segments)):
        segment = segments[i]
        if re.match(r'^[\u4e00-\u9fa5，。！？；：、——……（）]+$', segment):
            segments[i] = "[ZH]{}[ZH]".format(segment)
        elif re.match(r'^[a-zA-Z,.:()]+$', segment):
            segments[i] = "[EN]{}[EN]".format(segment)
        elif re.match(r'^\d+$', segment):
            segments[i] = "[ZH]{}[ZH]".format(segment)  # 数字视为中文
        else:
            segments[i] = "[JA]{}[JA]".format(segment)  # 日文

    return ''.join(segments)


def get_text(text, hps):
    text_cleaners = ['my_infer_ce_cleaners']
    text_norm = text_to_sequence(text, hps.symbols, text_cleaners)
    if hps.data.add_blank:
        text_norm = commons.intersperse(text_norm, 0)
    text_norm = torch.LongTensor(text_norm)
    return text_norm


def infer_one_audio(text, speaker_id=94, length_scale=1):
    '''
        input_type: 1输入自带语言标签  2中文  3中英混合
        length_scale: 语速，越小语速越快
    '''
    with torch.no_grad():
        stn_tst = get_text(text, hps)
        x_tst = stn_tst.to(device).unsqueeze(0)
        x_tst_lengths = torch.LongTensor([stn_tst.size(0)]).to(device)
        sid = torch.LongTensor([speaker_id]).to(device)  # speaker id
        audio = \
            net_g.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=.667, noise_scale_w=0.8, length_scale=length_scale)[
                0][0, 0].data.cpu().float().numpy()
        return audio
    return None


def infer_one_wav(text, speaker_id=94, length_scale=1, wav_name='1.wav'):
    '''
        input_type: 1输入自带语言标签  2中文  3中英混合
        length_scale: 语速，越小语速越快
    '''
    audio = infer_one_audio(text, speaker_id, length_scale)
    write('./infer_output/{}'.format(wav_name), hps.data.sampling_rate, audio)
    print('task done!')


def add_slience(wav_path, slience_len=100):
    silence = AudioSegment.silent(duration=slience_len)
    wav_audio = AudioSegment.from_wav(wav_path)
    wav_audio = wav_audio + silence
    wav_audio.export(wav_path, format="wav")
    pass


if __name__ == '__main__':
    infer_one_wav(
        '觉得本教程对你有帮助的话，记得一键三连哦！',
        speaker_id=0,
        length_scale=1.2)
    pass
